# encoding=utf-8
"""
需要这两个环境变量, 在pycharm命令行启动前设置
source /usr/local/Ascend/ascend-toolkit/set_env.sh
export PYTHONPATH=/usr/local/Ascend/thirdpart/aarch64/acllite:$PYTHONPATH
"""
import time
import cv2
import numpy
from AclTools import AclLiteResource
from acllite_model import AclLiteModel


class FundsStillGANAclInfer:
    def __init__(self, pth_path: str):
        self.acl_resource = AclLiteResource()
        self.acl_resource.init()
        self.model = AclLiteModel(pth_path)

    def softmax(self, x):
        f_x = numpy.exp(x) / numpy.sum(numpy.exp(x))
        return f_x

    def infer(self, image_path: str):
        trans_dim, ori_size = self._preprocess(image_path)
        result = self.model.execute([trans_dim, ])
        return self._postprocess(result[0][0], ori_size)

    def _postprocess(self, pred_arr: numpy.ndarray, ori_size: tuple):
        # ori size param
        H, W = ori_size
        # color channel
        image = (numpy.transpose(pred_arr, (1, 2, 0)) + 1) / 2.0 * 255.0
        # vis
        (r, g, b) = cv2.split(image)
        img_arr = cv2.merge([b, g, r])
        img_arr = cv2.resize(img_arr, (W, H), interpolation=cv2.INTER_CUBIC)
        img_bytes = self.array2bytes(img_arr, "png")
        image = numpy.asarray(bytearray(img_bytes), dtype=numpy.uint8)
        img_arr = cv2.imdecode(image, cv2.IMREAD_COLOR)
        return img_arr

    def _preprocess(self, data_in: str):
        # load
        ori_img = cv2.imread(data_in)
        ori_size = ori_img.shape[:2]  # h, w

        # transform
        trans_image = cv2.cvtColor(ori_img, cv2.COLOR_BGR2RGB)
        trans_image = cv2.resize(trans_image, (512, 512))
        trans_image = trans_image / 255.0
        img_arr = (trans_image - 0.5) / 0.5

        # return
        x = numpy.expand_dims(img_arr.astype(numpy.float32).transpose((2, 0, 1)), axis=0)
        return x, ori_size

    def array2bytes(self, array_img, suffix):
        # 对数组的图片格式进行编码
        success, encoded_array = cv2.imencode("." + suffix, array_img)
        # 将数组转为bytes
        bytes_img = encoded_array.tostring()

        return bytes_img


if __name__ == '__main__':
    # init
    t1 = time.perf_counter()
    enhance = FundsStillGANAclInfer("funds_om_enhance.om")
    t2 = time.perf_counter()
    print("init cost", t2 - t1)

    # time cost
    t1 = time.perf_counter()
    enhance_array = enhance.infer("1.jpg")
    t2 = time.perf_counter()
    print("infer cost", t2 - t1)
    cv2.imwrite("out.png", enhance_array)

    # query resource
    import os
    os.system("npu-smi info")

    # release
    del enhance
    time.sleep(5)
    os.system("npu-smi info")
